import subprocess
import time
import os
import signal

def auto_restart_pull(model_name: str, interval: int = 120):
    """
    每指定时间（默认120秒=2分钟）强制中断拉取，再重新拉取，直到成功。
    :param model_name: 要拉取的模型名（如 'qwen:7b-q4_0'）
    :param interval: 定时中断间隔（秒），默认2分钟
    """
    print(f"=== 开始定时中断拉取：模型={model_name}，间隔={interval}秒 ===")
    print(f"提示：每次中断后会自动从上次进度续传，拉取成功后会自动停止\n")

    while True:
        # 1. 启动 ollama pull 进程（记录进程ID，用于后续中断）
        pull_process = subprocess.Popen(
            ['ollama', 'pull', model_name],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            creationflags=subprocess.CREATE_NEW_PROCESS_GROUP  # Windows下允许强制杀死进程
        )
        process_id = pull_process.pid
        print(f"✅ 启动拉取进程（ID：{process_id}），将在{interval}秒后中断...")

        # 2. 等待指定时间（2分钟），同时监测拉取是否提前成功
        start_time = time.time()
        while True:
            time.sleep(5)  # 每5秒检查一次进程状态
            elapsed_time = time.time() - start_time

            # 检查拉取是否已经成功（进程结束且输出含 success）
            if pull_process.poll() is not None:  # 进程已结束
                stdout, stderr = pull_process.communicate()
                if 'success' in stdout:
                    print(f"\n🎉 模型拉取成功！进程{process_id}已完成")
                    return  # 成功，退出脚本
                else:
                    print(f"\n❌ 进程{process_id}提前结束（未成功），原因：{stderr[:150]}...")
                    break  # 重新启动拉取

            # 检查是否达到中断时间（2分钟）
            if elapsed_time >= interval:
                print(f"\n⏰ 已达到{interval}秒，强制中断进程{process_id}")
                # Windows下强制杀死进程（包括子进程）
                os.kill(process_id, signal.CTRL_BREAK_EVENT)
                pull_process.wait()  # 等待进程彻底退出
                print(f"🔚 进程{process_id}已中断，10秒后重新拉取...\n")
                time.sleep(10)  # 短暂间隔，避免频繁请求
                break

if __name__ == '__main__':
    # 替换为你需要拉取的模型（如 qwen:7b-q4_0 / llama2:7b-chat-q4_0）
    TARGET_MODEL = "qwen:7b-q4_0"
    # 定时中断间隔（秒），这里设为120秒=2分钟
    INTERVAL_SECONDS = 120
    # 启动脚本
    auto_restart_pull(TARGET_MODEL, INTERVAL_SECONDS)